#ifndef __MULTIGRID_H__
#define __MULTIGRID_H__

#include "Equation.h"
#include "Operator.h"
#include "Equation2D.h"
#include <ostream>
#include <fstream>
#include <utility>
#include <iomanip>

#define TEMPLATE template <int DIM, template <int> class Restriction, template <int> class Interpolation>
#define Template template <int DIM, template <int, template <int> class, template <int> class> class Cycle, template <int> class Restriction, template <int> class Interpolation>

template
<
	int DIM = 1,
	template <int> class Restriction = FullWeighting,
	template <int> class Interpolation = Linear
>
class VCycle
{
protected:
	int n;
	int v1;
	int v2;
	Vector boundary;
	const Vector SolveLayer(const Vector& initial_value, const Vector& f_h, int layer_n) const;

public:
	VCycle() = default;
	VCycle(int _n, const Vector& _boundary, int _v1 = 2, int _v2 = 2);
	virtual ~VCycle() = default;
	void SetN(int _n);
	int GetN() const;
	void SetV(int _v1, int _v2);
	void SetBoundary(const Vector& boundary);
	virtual const std::pair<Vector, double> Solve(const Vector& initial_value) const;
};

TEMPLATE
VCycle<DIM, Restriction, Interpolation>::VCycle(int _n, const Vector& _boundary, int _v1, int _v2) :
	n(_n), boundary(_boundary), v1(_v1), v2(_v2)
{
	if (DIM == 1 && _boundary.size() != 2)
		throw std::length_error("VCycle.SetBoundary(): boundary.size != 2");
	if (DIM == 2 && _boundary.size() != static_cast<size_t>(4 * n))
		throw std::length_error("VCycle.SetBoundary(): boundary.size != 4n");
}

TEMPLATE
void VCycle<DIM, Restriction, Interpolation>::SetN(int _n)
{
	n = _n;
}

TEMPLATE
int VCycle<DIM, Restriction, Interpolation>::GetN() const
{
	return n;
}

TEMPLATE
void VCycle<DIM, Restriction, Interpolation>::SetV(int _v1, int _v2)
{
	v1 = _v1;
	v2 = _v2;
}

TEMPLATE
void VCycle<DIM, Restriction, Interpolation>::SetBoundary(const Vector& _boundary)
{
	if (DIM == 1 && _boundary.size() != 2)
		throw std::length_error("VCycle.SetBoundary(): boundary.size != 2");
	if (DIM == 2 && _boundary.size() != static_cast<size_t>(4 * n))
		throw std::length_error("VCycle.SetBoundary(): boundary.size != 4n");
	boundary = _boundary;
}

TEMPLATE
const std::pair<Vector, double> VCycle<DIM, Restriction, Interpolation>::Solve(const Vector& initial_value) const
{
	if (DIM == 2) {
		if (initial_value.size() != static_cast<size_t>((n - 1)*(n - 1)))
			throw std::length_error("VCycle.Solve(): initial_value.size != (n-1)^2");
		PossionMatrix2D A(n);
		Vector rhs(RHS2D(n, boundary).GetVector());
		Vector v(SolveLayer(initial_value, rhs, n));
		return std::make_pair(v, MaxNorm(rhs - A * v));
	}
	else if (DIM == 1) {
		if (initial_value.size() != static_cast<size_t>(n - 1))
			throw std::length_error("VCycle.Solve(): initial_value.size != n-1");
		PossionMatrix A(n);
		Vector rhs(RHS(n, boundary[0], boundary[1]).GetVector());
		Vector v(SolveLayer(initial_value, rhs, n));
		return std::make_pair(v, MaxNorm(rhs - A * v));
	}	
}

TEMPLATE
const Vector VCycle<DIM, Restriction, Interpolation>::SolveLayer(const Vector& initial_value, const Vector& f_h, int layer_n) const
{
	if (DIM == 1 && initial_value.size() != static_cast<size_t>(layer_n - 1))
		throw std::length_error("VCycle.SolveLayer(): initial_value.size != layer_n-1");
	if (DIM == 2 && initial_value.size() != static_cast<size_t>((layer_n - 1) * (layer_n - 1)))
		throw std::length_error("VCycle.SolveLayer(): initial_value.size != (layer_n-1)^2");
	Vector v_h(initial_value), v_2h, f_2h;
	if (DIM == 1) {
		PossionMatrix A(layer_n);
		WeightedJacobi weighted_jacobi(A, f_h);
		for (int i = 0; i < v1; ++i)
			weighted_jacobi.IterateOnce(v_h);
		if (layer_n >= 4) {
			f_2h = Restriction<DIM>(layer_n) * (f_h - A * v_h);
			v_2h = Vector(layer_n / 2 - 1);
			v_2h = SolveLayer(v_2h, f_2h, layer_n / 2);
			v_h += Interpolation<DIM>(layer_n) * v_2h;
		}
		for (int i = 0; i < v2; ++i)
			weighted_jacobi.IterateOnce(v_h);
	}
	else if (DIM == 2) {
		PossionMatrix2D A(layer_n);
		WeightedJacobi2D weighted_jacobi(A, f_h);
		for (int i = 0; i < v1; ++i)
			weighted_jacobi.IterateOnce(v_h);
		if (layer_n >= 4) {
			f_2h = Restriction<DIM>(layer_n) * (f_h - A * v_h);
			v_2h = Vector((layer_n / 2 - 1) * (layer_n / 2 - 1));
			v_2h = SolveLayer(v_2h, f_2h, layer_n / 2);
			v_h += Interpolation<DIM>(layer_n) * v_2h;
		}
		for (int i = 0; i < v2; ++i)
			weighted_jacobi.IterateOnce(v_h);
	}

	return v_h;
}



template
<
	int DIM = 1,
	template <int> class Restriction = FullWeighting,
	template <int> class Interpolation = Linear
>
class FullMultigridVCycle : public VCycle<DIM, Restriction, Interpolation>
{
private:
	const Vector SolveLayer(const Vector& f_h, int layer_n) const;

public:
	FullMultigridVCycle() = default;
	FullMultigridVCycle(int _n, const Vector& _boundary, int _v1 = 2, int _v2 = 2);
	~FullMultigridVCycle() = default;
	virtual const std::pair<Vector, double> Solve(const Vector& initial_value = Vector(1)) const;
};

TEMPLATE
FullMultigridVCycle<DIM, Restriction, Interpolation>::FullMultigridVCycle(int _n, const Vector& _boundary, int _v1, int _v2) :
	VCycle<DIM, Restriction, Interpolation>(_n, _boundary, _v1, _v2) {}


TEMPLATE
const std::pair<Vector, double> FullMultigridVCycle<DIM, Restriction, Interpolation>::Solve(const Vector& initial_value) const
{
	if (DIM == 1) {
		PossionMatrix A(VCycle<DIM, Restriction, Interpolation>::n);
		Vector rhs(RHS(VCycle<DIM, Restriction, Interpolation>::n,
					VCycle<DIM, Restriction, Interpolation>::boundary[0],
					VCycle<DIM, Restriction, Interpolation>::boundary[1]).GetVector());
		Vector v(SolveLayer(rhs, VCycle<DIM, Restriction, Interpolation>::n));
		return std::make_pair(v, MaxNorm(rhs - A * v));
	}
	else if (DIM == 2) {
		PossionMatrix2D A(VCycle<DIM, Restriction, Interpolation>::n);
		Vector rhs(RHS2D(VCycle<DIM, Restriction, Interpolation>::n,
					VCycle<DIM, Restriction, Interpolation>::boundary).GetVector());
		Vector v(SolveLayer(rhs, VCycle<DIM, Restriction, Interpolation>::n));
		return std::make_pair(v, MaxNorm(rhs - A * v));
	}
	else
		return std::make_pair(Vector(1), -1);
}

TEMPLATE
const Vector FullMultigridVCycle<DIM, Restriction, Interpolation>::SolveLayer(const Vector& f_h, int layer_n) const
{
	Vector v_h, v_2h, f_2h;
	if (DIM == 1) {
		if (layer_n >= 4) {
			f_2h = Restriction<DIM>(layer_n) * f_h;
			v_2h = SolveLayer(f_2h, layer_n / 2);
			v_h = Interpolation<DIM>(layer_n) * v_2h;
		}
		else {
			v_h = Vector(layer_n - 1);
		}
		for (int i = 0; i < 10; ++i)
			v_h = VCycle<DIM, Restriction, Interpolation>::SolveLayer(v_h, f_h, layer_n);
	}
	else if (DIM == 2) {
		if (layer_n >= 4) {
			f_2h = Restriction<DIM>(layer_n) * f_h;
			v_2h = SolveLayer(f_2h, layer_n / 2);
			v_h = Interpolation<DIM>(layer_n) * v_2h;
		}
		else {
			v_h = Vector((layer_n - 1) * (layer_n - 1));
		}
		for (int i = 0; i < 10; ++i)
			v_h = VCycle<DIM, Restriction, Interpolation>::SolveLayer(v_h, f_h, layer_n);
	}
	
	return v_h;
}


// to help factory
class AbstractMultigrid
{
public:
	AbstractMultigrid() = default;
	virtual ~AbstractMultigrid() = default;
	virtual void SetParams(int _n, const Vector& _init_guess, const Vector& _boundary,
		int _max_iterations = 100, double _relative_accuracy = 0.01, int _v1 = 2, int _v2 = 2) = 0;
	virtual void SetInitialGuess(const Vector& v) = 0;
	virtual void SetMaxIterations(int _max_iterations) = 0;
	virtual void SetRelativeAccuracy(double _relative_accuracy) = 0;
	virtual const Vector Solve(std::ostream& os = std::cout) const = 0;
};

template
<
	int DIM = 1,
	template <int, template <int> class, template <int> class> class Cycle = VCycle,
	template <int> class Restriction = FullWeighting,
	template <int> class Interpolation = Linear
>
class Multigrid : public AbstractMultigrid
{
private:
	Cycle<DIM, Restriction, Interpolation> cycle;
	int max_iterations;
	double relative_accuracy;
	Vector init_guess;

public:
	Multigrid() = default;
	Multigrid(int _n, const Vector& _init_guess, const Vector& _boundary,
		int _max_iterations = 100, double _relative_accuracy = 1e-8, int _v1 = 2, int _v2 = 2);
	~Multigrid() = default;
	void SetParams(int _n, const Vector& _init_guess, const Vector& _boundary,
		int _max_iterations = 100, double _relative_accuracy = 1e-8, int _v1 = 2, int _v2 = 2);
	void SetInitialGuess(const Vector& v);
	void SetMaxIterations(int _max_iterations);
	void SetRelativeAccuracy(double _relative_accuracy);
	void SetCycle(const Cycle<DIM, Restriction, Interpolation>& _cycle);
	const Vector Solve(std::ostream& os = std::cout) const;
};

/*
Template
Multigrid<DIM, Cycle, Restriction, Interpolation>::Multigrid() :
	cycle(8, Vector(8), 2, 2), max_iterations(10),
	relative_accuracy(0.01), init_guess(Vector(8)) {}
*/

Template
Multigrid<DIM, Cycle, Restriction, Interpolation>::Multigrid(int _n, const Vector& _init_guess, const Vector& _boundary,
	int _max_iterations, double _relative_accuracy, int _v1, int _v2) :
	cycle(_n, _boundary, _v1, _v2), max_iterations(_max_iterations),
	relative_accuracy(_relative_accuracy), init_guess(_init_guess) {}

Template
void Multigrid<DIM, Cycle, Restriction, Interpolation>::SetParams(int _n, const Vector& _init_guess, const Vector& _boundary,
	int _max_iterations, double _relative_accuracy, int _v1, int _v2)
{
	cycle = Cycle<DIM, Restriction, Interpolation>(_n, _boundary, _v1, _v2);
	max_iterations = _max_iterations;
	relative_accuracy = _relative_accuracy;
	init_guess = _init_guess;
}

Template
void Multigrid<DIM, Cycle, Restriction, Interpolation>::SetInitialGuess(const Vector& v)
{
	init_guess = v;
}

Template
void Multigrid<DIM, Cycle, Restriction, Interpolation>::SetMaxIterations(int _max_iterations)
{
	max_iterations = _max_iterations;
}

Template
void Multigrid<DIM, Cycle, Restriction, Interpolation>::SetRelativeAccuracy(double _relative_accuracy)
{
	relative_accuracy = _relative_accuracy;
}

Template
void Multigrid<DIM, Cycle, Restriction, Interpolation>::SetCycle(const Cycle<DIM, Restriction, Interpolation>& _cycle)
{
	cycle = _cycle;
}

Template
const Vector Multigrid<DIM, Cycle, Restriction, Interpolation>::Solve(std::ostream& os) const
{
	int n = cycle.GetN();
	Vector v(init_guess), exact;
	double residual;
	if (DIM == 1) {
		exact = Vector(n - 1);
		for (int i = 0; i < n - 1; ++i)
			exact[i] = u((i + 1.0) * 1.0 / n);
	}
	else if (DIM == 2) {
		exact = Vector((n - 1) * (n - 1));
		for (int i = 1; i < n; ++i)
			for (int j = 1; j < n; ++j)
				exact[(n - 1) * (j - 1) + (i - 1)] = u2D(1.0 * i / n, 1.0 * j / n);
	}
	double exact_norm = MaxNorm(exact);
	double last_error, now_error;
	double last_residual, now_residual;
	for (int i = 0; i < max_iterations; ++i) {
		os << i + 1 << " ";
		std::tie(v, residual) = cycle.Solve(v);
		now_error = MaxNorm(v - exact);
		now_residual = residual;
		if (i > 0)
			os << now_error << " " << now_error / last_error
			   << " " << now_residual << " " << now_residual / last_residual << std::endl;
		else
			os << now_error << " 0 " << now_residual << " 0" << std::endl;
		last_error = now_error;
		last_residual = now_residual;
		if (MaxNorm(v - exact) / exact_norm < relative_accuracy)
			break;
	}

	return v;
}

#endif
